# -*- coding: utf-8 -*-

# @Time : 2022/3/18 9:42

# @Author : Aweo
# @File : CNN5.py

import torch
import torch.nn as nn

class CNN(nn.Module):
    def __init__(self, n_class):
        super(CNN, self).__init__()
        self.n_class = n_class  # 分类数

        # 卷积 + 激活 + 池化
        self.layer1 = nn.Sequential(
            nn.Conv2d(3, 6, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2, 2)
        )

        self.layer2 = nn.Sequential(
            nn.Conv2d(6, 12, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2, 2)
        )

        self.layer3 = nn.Sequential(
            nn.Conv2d(12, 24, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2, 2)
        )

        # 全连接层
        self.fc = nn.Sequential(
            nn.Linear(384, 128),
            nn.ReLU(),
            nn.Linear(128, self.n_class),
        )

    def forward(self, x):
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x

def cnn5():
    return CNN(n_class=2)